# Super Resolution model definition in PyTorch
import torch.nn as nn
import torch.onnx


# Load pretrained model weights
torch_model = torch.load("./best.pth")

# Set the model to inference mode
torch_model.eval()

# Export the model to ONNX
x = torch.randn(1, 3, 512, 512, requires_grad=True)
torch_out = torch_model(x)

torch.onnx.export(torch_model, x, "super_resolution.onnx", export_params=True, opset_version=10, do_constant_folding=True, input_names=['input'], output_names=['output'], dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}})
